// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <melon/cpp_attributes.h>
#include <melon/chrono/hardware.h>
#include <atomic>
#include <chrono>
#include <optional>

namespace kumo::pollux {
    /// Measures the time between construction and destruction with
    /// std::chrono::steady_clock and increments a user-supplied counter with the
    /// elapsed time in microseconds.
    class MicrosecondTimer {
    public:
        explicit MicrosecondTimer(uint64_t *timer) : timer_(timer) {
            start_ = std::chrono::steady_clock::now();
        }

        ~MicrosecondTimer() {
            auto duration = std::chrono::duration_cast<std::chrono::microseconds>(
                std::chrono::steady_clock::now() - start_);

            (*timer_) += duration.count();
        }

    private:
        std::chrono::steady_clock::time_point start_;
        uint64_t *timer_;
    };

    class NanosecondTimer {
    public:
        explicit NanosecondTimer(uint64_t *timer) : timer_(timer) {
            start_ = std::chrono::steady_clock::now();
        }

        ~NanosecondTimer() {
            auto duration = std::chrono::duration_cast<std::chrono::nanoseconds>(
                std::chrono::steady_clock::now() - start_);

            (*timer_) += duration.count();
        }

    private:
        std::chrono::steady_clock::time_point start_;
        uint64_t *timer_;
    };

    /// Measures the time between construction and destruction with CPU clock
    /// counter (rdtsc on X86) and increments a user-supplied counter with the cycle
    /// count.
    class ClockTimer {
    public:
        explicit ClockTimer(uint64_t &total)
            : total_(&total), start_(melon::hardware_timestamp()) {
        }

        explicit ClockTimer(std::atomic<uint64_t> &total)
            : atomicTotal_(&total), start_(melon::hardware_timestamp()) {
        }

        ~ClockTimer() {
            auto elapsed = melon::hardware_timestamp() - start_;
            if (total_) {
                *total_ += elapsed;
            } else {
                *atomicTotal_ += elapsed;
            }
        }

    private:
        uint64_t *total_{nullptr};
        std::atomic<uint64_t> *atomicTotal_{nullptr};
        uint64_t start_;
    };

    /// Returns the current epoch time in seconds.
    uint64_t getCurrentTimeSec();

    /// Returns the current epoch time in milliseconds.
    uint64_t getCurrentTimeMs();

    /// Returns the current epoch time in microseconds.
    uint64_t getCurrentTimeMicro();

    /// Returns the current epoch time in nanoseconds.
    uint64_t getCurrentTimeNano();
} // namespace kumo::pollux
